#include <gtest/gtest.h>

#include <future>
#include <vector>

#include "scheduler.hpp"
#include "task_set.hpp"

TEST(TaskSetTest, BinarySearch) {
  const int NUM = 1000;
  std::vector<int> data(NUM);
  for (int i = 0; i < NUM; ++i) {
    data[i] = i * 2;
  }
  const int target = 100;

  // define task func
  auto init = data.begin();
  auto search_task = [init](auto begin, auto end, int target) {
    std::vector<int> results;
    auto it = std::lower_bound(begin, end, target);

    while (it != end && *it == target) {
      results.push_back(std::distance(init, it));
      ++it;
    }
    return results;
  };

  // define final func
  auto finalizer = [](std::vector<std::vector<int>> &&partial) {
    auto part = std::move(partial);

    std::vector<int> final_result;
    for (auto &vec : part) {
      final_result.insert(final_result.end(), vec.begin(), vec.end());
    }
    return final_result;
  };

  // make task set and resource group
  auto task_set =
      make_task_set(data.begin(), data.end(), search_task, finalizer, target);
  task_set->_init_task_set(4);
  // use old scheduler to execute
  OldScheduler scheduler;
  auto future = task_set->get_future();
  while (true) {
    auto status = future.wait_for(std::chrono::milliseconds(10));  // NOLINT
    if (status == std::future_status::ready) {
      auto final_result = future.get();
      EXPECT_EQ(final_result.size(), 1);
      EXPECT_EQ(final_result[0], 50);
      return;
    }
    scheduler.submit_task(0,
                          [&task_set]() { return task_set->execute_task(0); });
    scheduler.submit_task(1,
                          [&task_set]() { return task_set->execute_task(1); });
    scheduler.submit_task(2,
                          [&task_set]() { return task_set->execute_task(2); });
    scheduler.submit_task(3,
                          [&task_set]() { return task_set->execute_task(3); });
  }
}
